package com.wfh.aigenerator.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.function.Consumer;

/**
 * @Title: SseService
 * @Author wangfenghuan
 * @Package com.wfh.aigenerator.service
 * @Date 2025/8/13 21:07
 * @description: SSE 服务，管理 SseEmitter 连接
 */
@Service
@Slf4j
public class SseService {

	// 存储每一个任务的 emitter
	private final Map<Long, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();
	// 存储尚未发送的事件（当客户端未连接时暂存）
	private final Map<Long, Queue<PendingEvent>> pendingEventQueueMap = new ConcurrentHashMap<>();
	private static final int MAX_PENDING_EVENTS_PER_TASK = 200;

	/**
	 * 创建一个新的 SSE 连接
	 *
	 * @param taskId 任务 ID
	 * @param timeout 超时时间（毫秒）
	 * @return SseEmitter
	 */
	public SseEmitter connect(Long taskId, Long timeout) {
		// 1. 先检查是否已有连接，有则关闭
		SseEmitter oldEmitter = sseEmitterMap.get(taskId);
		if (oldEmitter != null) {
			oldEmitter.complete();
			log.warn("关闭已存在的 SSE 连接，任务ID: {}", taskId);
		}

		// 2. 创建新的 emitter，设置超时时间
		SseEmitter sseEmitter = new SseEmitter(timeout); // ✅ 设置超时

		// 3. 添加回调，确保资源释放
		sseEmitter.onCompletion(() -> {
			log.info("SSE 连接完成，任务ID: {}", taskId);
			sseEmitterMap.remove(taskId);
			pendingEventQueueMap.remove(taskId);
		});

		sseEmitter.onTimeout(() -> {
			log.warn("SSE 连接超时，任务ID: {}", taskId);
			sseEmitter.complete(); // ✅ 主动关闭
			sseEmitterMap.remove(taskId);
			pendingEventQueueMap.remove(taskId);
		});

		sseEmitter.onError((throwable) -> {
			log.error("SSE 连接发生错误，任务ID: {}", taskId, throwable);
			sseEmitter.completeWithError(throwable); // ✅ 主动关闭
			sseEmitterMap.remove(taskId);
			pendingEventQueueMap.remove(taskId);
		});

		// 4. 存入 map
		sseEmitterMap.put(taskId, sseEmitter);
		log.info("创建新的 SSE 连接，taskId: {}, timeout: {}ms", taskId, timeout);

		// 5. 若存在暂存事件，则在连接建立后立即补发
		Queue<PendingEvent> queue = pendingEventQueueMap.remove(taskId);
		if (queue != null && !queue.isEmpty()) {
			log.info("为任务 {} 回放 {} 条暂存事件", taskId, queue.size());
			PendingEvent pending;
			while ((pending = queue.poll()) != null) {
				try {
					sseEmitter.send(SseEmitter.event().name(pending.eventName).data(pending.data));
				} catch (IOException e) {
					log.error("回放暂存事件失败，taskId: {}, event: {}", taskId, pending.eventName, e);
					sseEmitter.completeWithError(e);
					sseEmitterMap.remove(taskId);
					break;
				}
			}
		}

		return sseEmitter;
	}

	/**
	 * 发送消息给指定的任务
	 *
	 * @param taskId 任务 ID
	 * @param eventMsg 事件名（可选）
	 * @param data 数据
	 * @return 是否发送成功（若未连接则暂存并返回 true）
	 */
	public boolean sendMessage(Long taskId, String eventMsg, Object data) {
		SseEmitter sseEmitter = sseEmitterMap.get(taskId);
		if (sseEmitter == null) {
			// 未连接，暂存事件，待连接后补发
			Queue<PendingEvent> queue = pendingEventQueueMap.computeIfAbsent(taskId, k -> new ConcurrentLinkedQueue<>());
			queue.add(new PendingEvent(eventMsg, data));
			// 控制每个任务的最大暂存条数，超过则丢弃最早的
			while (queue.size() > MAX_PENDING_EVENTS_PER_TASK) {
				queue.poll();
			}
			log.warn("任务 {} 的连接不存在，消息已暂存，event: {}", taskId, eventMsg);
			return true;
		}
		try {
			sseEmitter.send(SseEmitter.event().data(data).name(eventMsg));
			return true;
		} catch (IOException e) {
			log.error("SSE 发送消息失败，taskId: {}, eventMsg: {}, data: {}", taskId, eventMsg, data, e);
			sseEmitter.completeWithError(e);
			sseEmitterMap.remove(taskId);
			pendingEventQueueMap.remove(taskId);
			return false;
		}
	}

	/**
	 * 关闭某个连接
	 *
	 * @param taskId 任务 ID
	 */
	public void complete(Long taskId) {
		SseEmitter emitter = sseEmitterMap.get(taskId);
		if (emitter != null) {
			emitter.complete();
			sseEmitterMap.remove(taskId);
			pendingEventQueueMap.remove(taskId);
		}
	}

	/**
	 * 批量操作（可选）
	 */
	public void forEach(Consumer<SseEmitter> action) {
		sseEmitterMap.values().forEach(action);
	}

	/**
	 * 暂存的事件模型
	 */
	private static final class PendingEvent {
		private final String eventName;
		private final Object data;

		private PendingEvent(String eventName, Object data) {
			this.eventName = eventName;
			this.data = data;
		}
	}
}